'''
Created on 01/06/2010

@author: Glazner
'''
from multiprocessing.managers import BaseManager
from threading import Thread
import thread
import time
import socket
import logging
import traceback
from wp.GLOBALS import *

class RemoteManager(BaseManager): pass
RemoteManager.register('get_worker')

def threaded(f):
    def theard_dec(*args):
        t = Thread(target=f,args=args)
        t.start()
        return t
    return theard_dec

def unbreakable(f):
    def unbreakable_dec(*args,**kw):
        try:
            return f(*args,**kw)
        except Exception as e:
            logging.error("\n".join(traceback.format_stack()))
            logging.error(str(e))
    return unbreakable_dec

class Task(object):
    def __init__(self,func,args,callback,group):
        self.func = func
        self.args = args
        self.callback = callback
        self.group = group
        self.res = None
        
    @threaded
    def finish(self,res):
        self.res = res
        self.callback(res)
        
    def __repr__(self):
        return "<Task func={0},args={1},group={2} result={3}>".format(
                  self.func,self.args,self.group,self.res)

        
class RWorker(object):
    def __init__(self,host,pool,key):
        self.host = host
        self.pool = pool
        self.key = key
        self._lock = thread.allocate_lock()
        self._worker = None
        self._hired = False
    @threaded
    def connect(self):
        while 1:
            try:
                if self.pool.closed: break
                mgr = RemoteManager(address=(self.hostname,self.port),
                                     authkey=self.key)
                mgr.connect()
                self._worker = mgr.get_worker()
                self.pool.addActiveWorker(self)
                break
            except (socket.error,IOError) as e:
                logging.info(" ".join(map(str,(e,(self.hostname,self.port)))))
                time.sleep(4)
    @property
    def port(self):
        if ":" in self.host:
            return int(self.host.split(":")[1])
        return DEFAULT_PORT
    
    @property
    def hostname(self):
        return self.host.split(":")[0]
        
                
    
    def execute(self,*args,**kw):
        return self._worker.execute_remote_task(*args,**kw)
        
    def hire(self):
        with self._lock:
            if not self._hired: 
                self._hired = True
                return self._hired
            return False
        
    def fire(self):
        with self._lock:
            assert self._hired
            self._hired = False

class WorkerPoolEvents(object):
    '''Interface for recieving callbacks from WorkerPool'''
    def __getattribute__(self,name):
        attr = object.__getattribute__(self,name)
        if name.startswith("on"):
            return unbreakable(attr)
        return attr
    
    def onHostConnected(self,host):
        pass
    
    def onHostConnectionLost(self,host,exception):
        pass
    
    def onHostIdle(self,host):
        pass
    
    def onHostStartWorking(self,host,task):
        pass

class WorkerPool(object):
    def __init__(self,hosts,poolEvents=WorkerPoolEvents(),key=DEFAULT_KEY):
        self._lock = thread.allocate_lock()
        self.hosts = tuple(hosts)
        self._poolEvents = poolEvents
        self._workers = {}
        self._active_workers = {}
        self._waittasks = []
        self._queue = []
        self._closed = False
        for host in hosts:
            self._workers[host] = RWorker(host,self,key) 
        for w in self._workers.values(): w.connect()
        
    def addActiveWorker(self,rworker):
        with self._lock:
            if self._closed: return
            self._active_workers[rworker.host]=rworker
            self._poolEvents.onHostConnected(rworker.host)
        self._schedule()
        
    def submit(self,f,args,callback=lambda *args:None,group='default'):
        task = Task(f,args,callback,group)
        with self._lock:
            self._waittasks.append(task)
            self._queue.append(task)
        self._schedule()
        return task
            
    def _schedule(self):
        with self._lock:
            while self._queue:
                task=self._queue.pop(0)
                for host,rworker in self._active_workers.items():
                    if rworker.hire():
                        self._run(rworker,task)
                        break
                else:
                    self._queue.insert(0, task)
                    break
    @threaded
    def _run(self,rworker,task):
        try:
            self._poolEvents.onHostStartWorking(rworker.host,task)
            res = rworker.execute(task.func,*task.args)
            task.finish(res)
            with self._lock:
                self._waittasks.remove(task)
            rworker.fire()
            self._poolEvents.onHostIdle(rworker.host)
            self._schedule()
        except IOError as e:
            rworker.fire()
            with self._lock:
                if task in self._waittasks:#maybe it was canceled
                    self._queue.insert(0,task)
                del self._active_workers[rworker.host]
            self._poolEvents.onHostConnectionLost(rworker.host,e)
            rworker.connect() #none blocking call
            
    def stop(self,group=None):
        with self._lock:
            toRemove = [t for t in self._queue if t.group == group]
            for t in toRemove:
                self._queue.remove(t)
                self._waittasks.remove(t)
    
    def wait(self,group=None):
        while True:
            with self._lock:
                if not self._waittasks:
                    break
                elif group:
                    if not [t for t in self._waittasks if t.group == group]:
                        break
            time.sleep(0.1)
            
    def getActiveNodes(self):
        with self._lock:
            return tuple(self._active_workers)
            
    def close(self):
        with self._lock:
            self._closed = True
    @property
    def closed(self):
        with self._lock: return self._closed

def alive_callback(alive): print 'got alive!'

def main():
    s = WorkerPool(['localhost'])
    for i in xrange(5):
        s.submit('alive',[],alive_callback)
    s.wait()
    
if __name__ == '__main__':
    main()
    